package org.marsdonne.auth;

import com.google.common.collect.Maps;
import org.springframework.security.jwt.JwtHelper;
import org.springframework.security.jwt.crypto.sign.RsaSigner;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.util.JsonParser;
import org.springframework.security.oauth2.common.util.JsonParserFactory;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

import java.security.KeyPair;
import java.security.interfaces.RSAPrivateKey;
import java.util.Map;

public class JwsAccessTokenConverter extends JwtAccessTokenConverter {

    private Map<String, String> headers;
    private JsonParser jsonParser = JsonParserFactory.create();
    final RsaSigner signer;

    public JwsAccessTokenConverter(Map<String, String> headers, KeyPair keyPair) {
        super();
        super.setKeyPair(keyPair);
        this.signer = new RsaSigner((RSAPrivateKey) keyPair.getPrivate());
        this.headers = headers;
    }

    @Override
    protected String encode(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        UserSecret user = (UserSecret) authentication.getPrincipal();

        Map<String, Object> content = Maps.newLinkedHashMap();
        content.put(User.Constants.USER_ID, user.getId());
        content.put(User.Constants.USER_REAL_NAME, user.getRealName());
        content.put(User.Constants.USED_CLIENT, user.getClient());
        content.put(User.Constants.USER_ORG_ID, user.getOrgId());
        content.put(User.Constants.USER_ORG_LICENSE_CODE, user.getOrgLicenseCode());
        content.put(User.Constants.USER_ORG_NAME, user.getOrgName());

        content.putAll(getAccessTokenConverter().convertAccessToken(accessToken, authentication));

        String source;

        try {
            source = this.jsonParser.formatMap(content);
        } catch (Exception ex) {
            throw new IllegalStateException("Cannot convert access token to JSON", ex);
        }

        return JwtHelper.encode(source, this.signer, this.headers).getEncoded();
    }
}